import csv

from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
from nltk.translate.chrf_score import sentence_chrf
from nltk.translate.gleu_score import sentence_gleu
from nltk.translate.nist_score import sentence_nist
from nltk.translate.meteor_score import meteor_score
from nltk.translate.ribes_score import sentence_ribes
from nltk import word_tokenize

from rich.progress import track

source_sentences = [
    "P13-7 X channel heading rudder lock drive high signal",
    "Crack at the position of the x bolt hole on the inner side of the lower edge strip of the actuator inner earpiece of the outer section of the right front wall",
    "Composite reinforcement",
    "Press-fit bushings",
    "Left/Right turn axis beam root after upper/lower side transition fillet R area",
    "The structure of a certain aircraft has been fully exposed to fatigue cracks, and the purpose of durability test has been achieved. After repair/reinforcement, the main structure has a significant difference from the field aircraft.",
    "Damage caused by design spectrum X FH is greater than that caused by field spectrum Y FH.",
    "The load spectrum (design spectrum) of the full-scale fatigue test is much more severe than the service load conditions of the field aircraft. A certain aircraft has exposed many fatigue cracks, and the purpose of the durability test has been basically achieved.",
    "Retain cracks in the right X frame, front wall of the wing, skin structure, etc., and only dispose of non-evaluated area cracks on the left side.",
    "Twice-cooked pork is only stir-fried when eating.",
    "The design spectrum flies heavy, while the load spectrum flies light.",
    "Damage tolerance test",
    "Durability test",
]

target_sentences = [
    "P13-7 Channel X roll motor electromagnetic lock driving signal (high-voltage end)",
    "Crack in bolt hole X of lower flange at inner section of inboard lug of right wing outboard front wall",
    "Reinforcement by composite materials",
    "Bush compressed by cold extrusion",
    "R zone on the upper/lower side of left/right pivot beam root (AFT)",
    "The aircraft has exposed a great many structual fatigue cracks, so the goal of durability test has been achieved; after multiple times of repair and reinforcement, the difference in the main structure status between aircraft under the design spectrum and those in the fleet is considerable.",
    "Damage under the design spectrum after X FH is greater than that under the fleet service spectrum after Y FH.",
    "The load spectrum (design spectrum) of the full-scale fatigue test is much more critical than the service load of the fleet. The aircraft has exposed a great number of fatigue cracks, basically achieving the objectives of the durability test.",
    "The cracks in the right wing frame X front wall and those in the skin structure shall be remained, and only the cracks in the left structure (zones not for test) shall be treated.",
    "We discussed this issue in previous meetings, so we do not need to go over that again.",
    "The load under the design spectrum is much more than that under the fleet service spectrum.",
    "Damage tolerance test",
    "Durability test",
]

bleu_scores = []
chrf_scores = []
gleu_scores = []
nist_scores = []
meteor_scores = []
ribes_scores = []

for i, (src, tgt) in track(enumerate(zip(source_sentences, target_sentences))):
    src_tokenized = word_tokenize(src)
    tgt_tokenized = word_tokenize(tgt)

    smooth = SmoothingFunction()
    bleu_score = sentence_bleu([src_tokenized], tgt_tokenized, smoothing_function=smooth.method1)
    try:
        nist_score = sentence_nist([src_tokenized], tgt_tokenized)
    except ZeroDivisionError:
        nist_score = 0
    chrf_score = sentence_chrf([src], tgt)
    gleu_score = sentence_gleu([src], tgt)
    meteor = meteor_score([src_tokenized], tgt_tokenized)
    ribes = sentence_ribes([src_tokenized], tgt_tokenized)

    bleu_scores.append(bleu_score)
    nist_scores.append(nist_score)
    chrf_scores.append(chrf_score)
    gleu_scores.append(gleu_score)
    meteor_scores.append(meteor)
    ribes_scores.append(ribes)

with open("results.csv", "w", encoding="utf-8", newline="") as csv_file:
    csv_writer = csv.writer(csv_file)

    names = ["meteor", "chrf", "gleu"]

    csv_writer.writerow(names)

    data = []
    for meteor, chrf_score, gleu_score in zip(
            meteor_scores,
            chrf_scores,
            gleu_scores):
        data.append([meteor, chrf_score, gleu_score])

    data.append([
        sum(meteor_scores)/len(meteor_scores),
        sum(chrf_scores)/len(chrf_scores),
        sum(gleu_scores)/len(gleu_scores)
    ])

    csv_writer.writerows(data)
